{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "outputs": [],
   "source": [
    "import nltk\n",
    "import string\n",
    "%matplotlib inline\n",
    "from nltk import word_tokenize\n",
    "import matplotlib.pyplot as plt\n",
    "from nltk.corpus import stopwords"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "outputs": [],
   "source": [
    "enstop = stopwords.words('english')\n",
    "punct = string.punctuation\n",
    "\n",
    "def tokenizer(sent):\n",
    "    sent = sent.lower()\n",
    "    tmp = word_tokenize(sent)\n",
    "    res = []\n",
    "    for word in tmp:\n",
    "        if word not in enstop and word not in punct:\n",
    "            res.append(word)\n",
    "    return res"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchtext import data\n",
    "from torchtext import vocab"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "outputs": [],
   "source": [
    "text_field = data.Field(tokenize=tokenizer, lower=True, include_lengths=True,\n",
    "                        fix_length=256)\n",
    "label_field = data.Field(sequential=False, use_vocab=False, dtype=torch.long)\n",
    "train, valid, test = data.TabularDataset.splits(path='',\n",
    "                                                train='Train.csv',\n",
    "                                                validation='Valid.csv',\n",
    "                                                test='Test.csv',\n",
    "                                                format='csv', skip_header=True,\n",
    "                                                fields=[('sentence', text_field), ('label', label_field)])"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "outputs": [],
   "source": [
    "vec = vocab.Vectors(name='glove.6B.300d.txt')\n",
    "text_field.build_vocab(train, valid, test, max_size=250000, vectors=vec,\n",
    "                       unk_init=torch.Tensor.normal_)\n",
    "label_field.build_vocab(train, valid, test)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "train_iter, valid_iter, test_iter = data.BucketIterator.splits((train, valid, test), batch_sizes=(64, 64, 64),\n",
    "                                                               sort_key=lambda x: len(x.sentence),\n",
    "                                                               sort_within_batch=True,\n",
    "                                                               repeat=False, shuffle=True,\n",
    "                                                               device=device)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "outputs": [],
   "source": [
    "def train_fun(model, train_iter, dev_iter, num_epoch, opt, criterion, eva,\n",
    "              out_model_file):\n",
    "    model.train()\n",
    "    loss_list = []\n",
    "    dev_acc = []\n",
    "    best_dev_acc = 0.\n",
    "    for epoch in range(num_epoch):\n",
    "        total_loss = 0.\n",
    "        for batch in train_iter:\n",
    "            output = model(batch.sentence)\n",
    "            loss = criterion(output, batch.label)\n",
    "            opt.zero_grad()\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "            total_loss += loss.item()\n",
    "        loss_list.append(total_loss)\n",
    "        dev_acc.append(eva(model, dev_iter))\n",
    "        print(f\"Epoch: {epoch+1}/{num_epoch}. Total loss: {total_loss:.3f}. Validation Set Acc: {dev_acc[-1]:.3%}.\")\n",
    "        if dev_acc[-1] > best_dev_acc:\n",
    "            best_dev_acc = dev_acc[-1]\n",
    "            torch.save(model.state_dict(), out_model_file)\n",
    "    return loss_list, dev_acc"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "outputs": [],
   "source": [
    "def eva(model, data_iter):\n",
    "    correct, count = 0, 0\n",
    "    with torch.no_grad():\n",
    "        for batch in data_iter:\n",
    "            pred = model(batch.sentence)\n",
    "            pred = torch.argmax(pred, dim=-1)#biaoqian\n",
    "            correct += (pred == batch.label).sum().item()\n",
    "            count += len(pred)\n",
    "    return correct / count"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "outputs": [],
   "source": [
    "class LSTM(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, num_layers, bidirectional,\n",
    "                 dropout_rate):\n",
    "        super(LSTM, self).__init__()\n",
    "        self.rnn = nn.LSTM(input_size=input_size, hidden_size=hidden_size,\n",
    "                           num_layers=num_layers, bidirectional=bidirectional,\n",
    "                           dropout=dropout_rate)\n",
    "\n",
    "    def forward(self, x, length):\n",
    "        packed_x = nn.utils.rnn.pack_padded_sequence(x, length)\n",
    "        packed_output, (hidden, cell) = self.rnn(packed_x)\n",
    "        output, output_lengths = nn.utils.rnn.pad_packed_sequence(packed_output)\n",
    "        return hidden, output"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "outputs": [],
   "source": [
    "class GRU(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, num_layers, bidirectional,\n",
    "                dropout_rate):\n",
    "        super(GRU, self).__init__()\n",
    "        self.rnn = nn.GRU(input_size=input_size, hidden_size=hidden_size,\n",
    "                          num_layers=num_layers, bidirectional=bidirectional,\n",
    "                          dropout=dropout_rate)\n",
    "\n",
    "    def forward(self, x, length):\n",
    "        packed_x = nn.utils.rnn.pack_padded_sequence(x, length)\n",
    "        packed_output, hidden = self.rnn(packed_x)\n",
    "        output, output_lengths = nn.utils.rnn.pad_packed_sequence(packed_output)\n",
    "        return hidden, output"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "outputs": [],
   "source": [
    "class TextRNN(nn.Module):\n",
    "    def __init__(self, embed_size, hidden_size, num_layers, bidirectional, out_dim,\n",
    "                 dropout_rate, pretrained_embed, use_gru=False, freeze=True,\n",
    "                 random_embed=False, vocab_size=None):\n",
    "        super(TextRNN, self).__init__()\n",
    "        if random_embed:\n",
    "            self.embed = nn.Embedding(vocab_size, embed_size)\n",
    "        else:\n",
    "            self.embed = nn.Embedding.from_pretrained(pretrained_embed, freeze=True)\n",
    "        if use_gru:\n",
    "            self.rnn = GRU(embed_size, hidden_size, num_layers, bidirectional,\n",
    "                           dropout_rate)\n",
    "        else:\n",
    "            self.rnn = LSTM(embed_size, hidden_size, num_layers, bidirectional,\n",
    "                            dropout_rate)\n",
    "        self.proj = nn.Linear(2*hidden_size, out_dim)\n",
    "        self.dropout = nn.Dropout(p=dropout_rate)\n",
    "\n",
    "    def forward(self, x):\n",
    "        text, text_length = x # text: [seq_len, bs]\n",
    "        text = text.permute(1, 0) # text: [bs, seq_len]\n",
    "        embed_x = self.embed(text) # embed_x: [bs, seq_len, embed_dim]\n",
    "        embed_x = embed_x.permute(1, 0, 2) # embed_x: [seq_len, bs, embed_dim]\n",
    "        hidden, _ = self.rnn(embed_x, text_length) # hidden: [2*num_layers, bs, hidden_size]\n",
    "        hidden = torch.cat((hidden[-1,:,:], hidden[-2,:,:]), dim=1)\n",
    "        return self.proj(self.dropout(hidden))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training begin!\n"
     ]
    }
   ],
   "source": [
    "embed_size = 300\n",
    "hidden_size = 300\n",
    "num_layers = 2\n",
    "bidirectional = True\n",
    "out_dim = 2\n",
    "dropout_rate = 0.2\n",
    "pretrained_embed = text_field.vocab.vectors\n",
    "lr = 0.001\n",
    "num_epoch = 20\n",
    "freeze = True\n",
    "use_gru = True\n",
    "random_embed = False\n",
    "vocab_size = len(text_field.vocab.stoi)\n",
    "out_model_file = 'textrnn_gru_freeze.pt'\n",
    "\n",
    "textrnn_gru_freeze = TextRNN(embed_size, hidden_size, num_layers, bidirectional, out_dim,\n",
    "                             dropout_rate, pretrained_embed, use_gru=use_gru, freeze=freeze,\n",
    "                             random_embed=random_embed, vocab_size=None).to(device)\n",
    "opt = torch.optim.Adam(textrnn_gru_freeze.parameters(), lr=lr)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "print(\"Training begin!\")\n",
    "loss_list, dev_acc_list = train_fun(textrnn_gru_freeze, train_iter, valid_iter, num_epoch, opt, criterion,\n",
    "                                    eva, out_model_file)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15,5))\n",
    "plt.plot(loss_list, marker='s')\n",
    "plt.xlabel('epoch')\n",
    "plt.ylabel('total loss')\n",
    "plt.title('Total loss per epoch')\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15,5))\n",
    "plt.plot(dev_acc_list, marker='s')\n",
    "plt.xlabel('epoch')\n",
    "plt.ylabel('dev acc')\n",
    "plt.title('Validation set accuracy')\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "textrnn_lstm_random = TextRNN(embed_size, hidden_size, num_layers, bidirectional, out_dim,\n",
    "                              dropout_rate, pretrained_embed, use_gru=use_gru, freeze=freeze,\n",
    "                              random_embed=random_embed, vocab_size=vocab_size).to(device)\n",
    "textrnn_lstm_random.load_state_dict(torch.load('textrnn_lstm_random.pt'))\n",
    "print(f\"Test set acc: {eva(textrnn_lstm_random, test_iter):.2%}\")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}